#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
The MIT License (MIT)
Copyright © 2022 Walkline Wang (https://walkline.wang)
Gitee: https://gitee.com/walkline/micropython-resource-manager

Append user resource files to specified partition of MicroPython firmware
"""
import os
import sys
from pathlib import Path
from config import Config

try:
	from littlefs import LittleFS, LittleFSError
except ImportError:
	print(
'''
Please install littlefs-python first!

    pip3 install littlefs-python'''
	)
	exit(0)


LITTLEFS_IMAGE_TEMP = 'littlefs.img'
PARTITION_SIZE = Config.Partition.SIZE
BLOCK_SIZE = 4096
PROG_SIZE = 256
DISK_VERSION = 0x0002_0000 # 2.0

def show_help(self):
	print(
f'''
Usage:
    python3 {self} <firmware_file> <user_file_folder>''')

# https://blog.csdn.net/vocalsir/article/details/130300910
def convert_file_size(size) -> str:
	units = 'Bytes', 'KB', 'MB', 'GB', 'TB'
	unit = units[0]

	for i in range(1, len(units)):
		if size >= 1024:
			size /= 1024
			unit = units[i]
		else:
			break

	return f'{size:.2f} {unit}'

def __is_file_folder_exists(name, is_folder=False) -> bool:
	if is_folder:
		return os.path.exists(name) and os.path.isdir(name)
	else:
		return os.path.exists(name) and os.path.isfile(name)

# https://github.com/jrast/littlefs-python/blob/master/examples/mkfsimg.py
def __generate_littlefs_partition(folder):
	block_count = PARTITION_SIZE // BLOCK_SIZE

	if block_count * BLOCK_SIZE != PARTITION_SIZE:
		print('Failed, image size should be a multiple of block size.')
		exit(1)

	fs = LittleFS(
		block_size=BLOCK_SIZE,
		block_count=block_count,
		prog_size=PROG_SIZE,
		disk_version=DISK_VERSION,
	)

	folder_path = Path(folder)
	print(f'Packing files from {folder}:')

	for filename in folder_path.rglob('*'):
		lfs_filename = f'/{filename.relative_to(folder_path).as_posix()}'

		if filename.is_file():
			with open(filename, 'rb') as src_file:
				print(f'- Adding {lfs_filename} ({convert_file_size(filename.stat().st_size)})')

				try:
					with fs.open(lfs_filename, 'wb') as lfs_file:
						lfs_file.write(src_file.read())
				except LittleFSError as lfse:
					if lfse.code == LittleFSError.Error.LFS_ERR_NOSPC:
						print('\nFailed, no enough space left.')
						exit(1)
					else:
						print(f'Failed, {lfse}')
		elif filename.is_dir():
			fs.mkdir(lfs_filename)

	with open(LITTLEFS_IMAGE_TEMP, 'wb') as image_file:
		image_file.write(fs.context.buffer)

def main(argv):
	firmware, folder = argv

	if __is_file_folder_exists(firmware) and __is_file_folder_exists(folder, True):
		__generate_littlefs_partition(folder)

		file, ext = os.path.splitext(firmware)

		with open(firmware, 'rb') as input_file,\
			open(f'{file}_output{ext}', 'wb') as output_file,\
			open(LITTLEFS_IMAGE_TEMP, 'rb') as partition_file:

			# 复制固件内容到输出文件并移动指针到资源分区起始处
			output_file.write(input_file.read())
			output_file.seek(Config.Partition.OFFSET)

			if Config.Partition.IS_ESP32_ESP32S2:
				output_file.seek(-0x1000, 1)

			# 写入资源分区文件
			output_file.write(partition_file.read())
			output_file.flush()

		try:
			os.remove(LITTLEFS_IMAGE_TEMP)
		except:
			pass

		print(f'\nCompleted, output file: {file}_output{ext}')
	else:
		print(f'\nFailed, firmware_file or user_file_folder not exists.')


if __name__ == '__main__':
	argv = sys.argv

	if len(argv) in (1, 2) or len(argv) > 3:
		show_help(argv[0])
		exit(0)

	main(sys.argv[1:])
